from utils import *
from PCA import PCA
import torch
import numpy as np
import time
import pickle
import os
from skimage.metrics import peak_signal_noise_ratio


raw_data = preprocess()


def get_psnr_pr_1(filename: str):
    with open(filename, '+rb') as f:
        recon = pickle.load(f)
    data = raw_data
    result = np.array([peak_signal_noise_ratio(data[i], recon[i], data_range=256) for i in range(data.shape[0])])
    return result.mean(-1)


def draw_pr_1(filename: str, n_top: int):
    with open(filename, '+rb') as f:
        recon = pickle.load(f)[:100]
    show_with_grid(f'results/recovered_faces_top_{n_top}.jpg', 
                   recon.reshape(recon.shape[0], 32, 32).transpose(0, 2, 1), 10, 10)


def main_pr1(n_components: list = [10, 50, 100, 150], bins_valid=False, img_valid=False):
    """问题二的小题 (a)

    Args:
        n_components (list, optional): 分别降维到列表中的各个维度。默认为 [10, 50, 100, 150]。
        bins_valid (bool, optional): 重建图像是否已打包好。默认为否。
    """
    if not os.path.exists('AfterPCA'):
        os.makedirs('AfterPCA')
    if not os.path.exists('results'):
        os.makedirs('results')
    if not bins_valid:
        for n_component in n_components:
            data = torch.from_numpy(raw_data)
            pca = PCA(n_component=n_component, device='cpu')
            pca.fit(data)
            after_pca = pca.transform(data)
            print('after_pca.shape', after_pca.shape)
            recon = pca.reconstruct(after_pca)
            print('recon.shape', recon.shape)
            recon = recon.cpu().numpy()

            with open(f'AfterPCA/pr1_top_{n_component}.bin', '+wb') as f:
                pickle.dump(recon, f)

    if not img_valid:
        for n_top in n_components:
            draw_pr_1(f'AfterPCA/pr1_top_{n_top}.bin', n_top)

    with open('results/log_pr1.txt', '+w') as f:
        for n_top in n_components:
            psnr = get_psnr_pr_1(f'AfterPCA/pr1_top_{n_top}.bin')
            f.write(f'avg_psnr_top_{n_top}: {psnr}\n')


if __name__ == '__main__':
    begin = time.time()
    # show_with_grid(f'results/origin_faces.jpg', 
    #                raw_data.reshape(raw_data.shape[0], 32, 32).transpose(0, 2, 1), 10, 10)
    with torch.no_grad():
        main_pr1([10, 50, 100, 150], bins_valid=False, img_valid=False)  # 若重建图像已打好包，可把 bins_valid 设为真    # 重建图片同理
    end = time.time()
    print('Total time:', end - begin)